# -*- coding: utf-8 -*-
"""
思维链模板测试

这个测试模块专门用于测试经典逻辑推理段子的思维链处理能力。
通过一系列相互关联的脑筋急转弯问题，验证LLM是否能够：

1. 理解问题之间的逻辑关系
2. 保持上下文的连贯性
3. 生成符合预期的标准答案
4. 展现逻辑推理能力

测试场景说明：
这是一个经典的四连环脑筋急转弯，每个问题的答案都依赖于前面问题的结果：
- 大象装冰箱 → 长颈鹿装冰箱（需要先取出大象）
- 长颈鹿在冰箱 → 动物园运动会缺席
- 鳄鱼参加运动会 → 河里安全通过

测试方法：
1. 直接答案引导：提供标准答案格式，测试复述能力
2. 学习模式：让AI学习完整的问答对，测试理解能力
3. 上下文推理：测试基于前文的逻辑推理能力

评估指标：
- 答案准确性：是否符合标准答案
- 逻辑连贯性：是否理解问题间的关联
- 创意保持：是否保持了幽默和创意元素

@Author  : gaoyu
@Time    : 2025/7/22 17:05
@Function: 经典逻辑推理段子 - 强制标准答案版本

Examples:
    运行测试：
    $ python tests/思维链模板测试.py
    
    预期输出：
    - 每个问题的标准答案
    - AI学习后的复述结果
    - 逻辑推理的准确性评估

Note:
    - 确保Ollama服务在指定地址运行
    - 确保模型已正确加载
    - 测试结果可能因模型而异
"""
from typing import Dict, List, Any, Optional
from langchain.prompts import PromptTemplate
from langchain_community.llms.ollama import Ollama
from langchain_core.output_parsers import StrOutputParser
from loguru import logger

# ==================== 模型配置 ====================

# 初始化 Ollama 客户端
# 使用DeepSeek-R1模型，这是一个专门优化过的推理模型
# 在逻辑推理任务上表现优异
ollama = Ollama(
    model="deepseek-r1:32b",           # 32B参数的DeepSeek-R1模型
    base_url="http://10.48.0.81:11434"  # Ollama服务地址
)

# ==================== 提示模板设计 ====================

# 直接答案引导模板
# 这个模板的设计理念是通过明确的指导来获得标准答案
# 适用于需要精确复述的场景
direct_template = PromptTemplate(
    input_variables=["question", "expected_answer"],
    template=(
        "这是一个经典的脑筋急转弯段子，有标准答案。\n\n"
        "问题：{question}\n\n"
        "标准答案应该是：{expected_answer}\n\n"
        "请直接输出这个标准答案，不要添加额外解释："
    )
)

# ==================== 处理链构建 ====================

# 创建LangChain处理链
# 链的组成：提示模板 → LLM → 输出解析器
# 这种设计模式便于模块化和测试
chain = direct_template | ollama | StrOutputParser()

# ==================== 测试数据定义 ====================

# 经典四连环脑筋急转弯标准问答对
# 这个数据集设计用于测试LLM的逻辑推理和上下文理解能力
# 每个问题都与前面的问题存在逻辑关联，形成完整的推理链
standard_qa: List[Dict[str, str]] = [
    {
        # 第一环：基础操作步骤
        # 测试：基本的步骤分解能力
        "question": "给大象装冰箱，总共分几步？",
        "answer": "三步：第一步打开冰箱门，第二步把大象装进去，第三步关上冰箱门。",
        "key": "elephant",
        "logic_type": "基础步骤分解",
        "dependency": None  # 无依赖，是起始问题
    },
    {
        # 第二环：基于前文的变体操作
        # 测试：上下文理解和逻辑推理
        "question": "给长颈鹿装冰箱，总共分几步？",
        "answer": "四步：第一步打开冰箱门，第二步把大象拿出来，第三步把长颈鹿装进去，第四步关上冰箱门。",
        "key": "giraffe",
        "logic_type": "上下文推理",
        "dependency": "elephant"  # 依赖第一个问题的结果
    },
    {
        # 第三环：状态推断
        # 测试：基于前文状态的逻辑推断
        "question": "动物园开运动会，谁没来？",
        "answer": "长颈鹿没来，因为长颈鹿在冰箱里。",
        "key": "zoo",
        "logic_type": "状态推断",
        "dependency": "giraffe"  # 依赖第二个问题的结果
    },
    {
        # 第四环：间接逻辑推理
        # 测试：多层逻辑关系的理解
        "question": "你过一条河，鳄鱼住在河里，你没被咬，为什么？",
        "answer": "因为鳄鱼去参加动物园运动会了。",
        "key": "river",
        "logic_type": "间接推理",
        "dependency": "zoo"  # 依赖第三个问题的结果
    }
]

# ==================== 测试执行部分 ====================

logger.info("=== 经典脑筋急转弯段子（标准答案版）===")
logger.info("测试目标：验证LLM的逻辑推理和上下文理解能力")
logger.info("")

# 存储测试结果，用于后续分析
results: Dict[str, Any] = {}

logger.info("\n方法：尝试让AI学习标准答案")
logger.info("=" * 40)

# ==================== 学习模式测试 ====================

# 学习模式提示模板
# 这个模板让AI先学习完整的问答对，然后复述
# 测试AI是否能理解和保持逻辑关系
learning_template = PromptTemplate(
    input_variables=["all_qa"],
    template=(
        "这是一个经典的脑筋急转弯段子，请学习以下标准答案：\n\n"
        "{all_qa}\n\n"
        "现在请你完整复述这四个问题和答案，保持原有的逻辑关系："
    )
)

# 创建学习模式处理链
learning_chain = learning_template | ollama | StrOutputParser()


def build_learning_material(qa_data: List[Dict[str, str]]) -> str:
    """
    构建学习材料
    
    将标准问答对格式化为适合AI学习的文本格式。
    包含问题、答案和逻辑关系说明。
    
    Args:
        qa_data: 标准问答数据列表
        
    Returns:
        str: 格式化的学习材料
        
    Examples:
        >>> material = build_learning_material(standard_qa)
        >>> print(material)
        1. 给大象装冰箱，总共分几步？
           答案：三步：第一步打开冰箱门...
    """
    learning_material = ""
    for i, item in enumerate(qa_data, 1):
        learning_material += f"{i}. {item['question']}\n"
        learning_material += f"   答案：{item['answer']}\n"
        learning_material += f"   逻辑类型：{item.get('logic_type', '未知')}\n"
        if item.get('dependency'):
            learning_material += f"   依赖关系：基于'{item['dependency']}'问题的结果\n"
        learning_material += "\n"
    return learning_material


def evaluate_ai_response(ai_response: str, standard_qa: List[Dict[str, str]]) -> Dict[str, Any]:
    """
    评估AI回答的质量
    
    从多个维度评估AI的回答：
    1. 完整性：是否包含所有问题和答案
    2. 准确性：答案是否与标准答案匹配
    3. 逻辑性：是否保持了问题间的逻辑关系
    
    Args:
        ai_response: AI的回答文本
        standard_qa: 标准问答数据
        
    Returns:
        Dict: 评估结果
    """
    evaluation = {
        "completeness": 0,  # 完整性得分
        "accuracy": 0,     # 准确性得分
        "logic": 0,        # 逻辑性得分
        "details": []      # 详细分析
    }
    
    # 检查是否包含所有关键词
    for qa in standard_qa:
        if qa['key'] in ai_response.lower():
            evaluation["completeness"] += 1
            evaluation["details"].append(f"包含关键词: {qa['key']}")
    
    # 计算得分百分比
    total_questions = len(standard_qa)
    evaluation["completeness"] = (evaluation["completeness"] / total_questions) * 100
    
    return evaluation


# 执行学习模式测试
try:
    logger.info("正在构建学习材料...")
    learning_material = build_learning_material(standard_qa)
    logger.info(f"学习材料长度: {len(learning_material)} 字符")
    
    logger.info("正在让AI学习标准答案...")
    ai_response = learning_chain.invoke({"all_qa": learning_material})
    
    logger.info("AI学习后的回答：")
    logger.info("-" * 50)
    logger.info(ai_response)
    logger.info("-" * 50)
    
    # 评估AI回答
    evaluation = evaluate_ai_response(ai_response, standard_qa)
    logger.info(f"\n评估结果：")
    logger.info(f"完整性得分: {evaluation['completeness']:.1f}%")
    logger.info(f"评估详情: {evaluation['details']}")
    
    # 保存结果
    results["learning_mode"] = {
        "response": ai_response,
        "evaluation": evaluation
    }
    
except Exception as e:
    logger.error(f"AI学习失败：{type(e).__name__}: {e}")
    logger.error("详细错误信息", exc_info=True)
    logger.info("将使用标准答案作为备选方案")
    
    # 记录失败信息
    results["learning_mode"] = {
        "error": str(e),
        "fallback": "使用标准答案"
    }


# ==================== 模型配置和扩展 ====================

def try_different_model(model_name: str = "llama2:13b", 
                       base_url: str = "http://10.48.0.81:11434") -> Optional[Ollama]:
    """
    尝试不同的模型配置
    
    这个函数提供了切换不同LLM模型的能力，
    便于比较不同模型在逻辑推理任务上的表现。
    
    Args:
        model_name: 要使用的模型名称
        base_url: Ollama服务地址
        
    Returns:
        Optional[Ollama]: 配置好的模型实例，失败时返回None
        
    Examples:
        >>> # 尝试使用Llama2模型
        >>> llama2 = try_different_model("llama2:13b")
        >>> if llama2:
        ...     # 使用新模型进行测试
        ...     pass
        
        >>> # 尝试使用其他模型
        >>> qwen = try_different_model("qwen:14b")
    """
    try:
        logger.info(f"正在尝试加载模型: {model_name}")
        other_llm = Ollama(model=model_name, base_url=base_url)
        
        # 简单测试模型是否可用
        test_response = other_llm.invoke("测试")
        logger.info(f"模型 {model_name} 加载成功")
        return other_llm
        
    except Exception as e:
        logger.error(f"模型 {model_name} 加载失败: {e}")
        return None


def run_comparative_test(models: List[str]) -> Dict[str, Any]:
    """
    运行多模型对比测试
    
    使用不同的模型对同一组问题进行测试，
    比较它们在逻辑推理任务上的表现差异。
    
    Args:
        models: 要测试的模型名称列表
        
    Returns:
        Dict: 各模型的测试结果对比
    """
    comparative_results = {}
    
    for model_name in models:
        logger.info(f"\n=== 测试模型: {model_name} ===")
        model = try_different_model(model_name)
        
        if model:
            # 使用新模型创建处理链
            test_chain = learning_template | model | StrOutputParser()
            
            try:
                response = test_chain.invoke({"all_qa": build_learning_material(standard_qa)})
                evaluation = evaluate_ai_response(response, standard_qa)
                
                comparative_results[model_name] = {
                    "response": response,
                    "evaluation": evaluation,
                    "status": "success"
                }
                
                logger.info(f"模型 {model_name} 完整性得分: {evaluation['completeness']:.1f}%")
                
            except Exception as e:
                logger.error(f"模型 {model_name} 测试失败: {e}")
                comparative_results[model_name] = {
                    "error": str(e),
                    "status": "failed"
                }
        else:
            comparative_results[model_name] = {
                "error": "模型加载失败",
                "status": "unavailable"
            }
    
    return comparative_results


# 提示用户可以进行模型对比测试
logger.info("\n=== 扩展功能提示 ===")
logger.info("如需测试其他模型，可以调用以下函数：")
logger.info("- try_different_model('模型名称')")
logger.info("- run_comparative_test(['模型1', '模型2', ...])")
logger.info("\n支持的模型示例：")
logger.info("- llama2:13b, qwen:14b, chatglm:6b, baichuan:13b")


# ==================== 测试总结 ====================

logger.info("\n=== 测试总结 ===")
logger.info(f"已完成 {len(results)} 项测试")
for test_name, result in results.items():
    if "evaluation" in result:
        logger.info(f"{test_name}: 完整性 {result['evaluation']['completeness']:.1f}%")
    elif "error" in result:
        logger.info(f"{test_name}: 测试失败 - {result['error']}")

logger.info("\n思维链测试完成！")
